package sdwebuiapi

import (
	"encoding/json"
	"fmt"
	"github.com/gogf/gf/frame/g"
	"io/ioutil"
	"path/filepath"
	"sdwebuiapi/utils"
	"strconv"
)

type Data struct {
	Prompt         string `json:"prompt"`
	NegativePrompt string `json:"negative_prompt"`
	Seed           int    `json:"seed"`
	//SamplerName      string  `json:"sampler_name"`
	SamplerIndex     string  `json:"sampler_index"`
	CfgScale         float64 `json:"cfg_scale"`
	Width            int     `json:"width"`
	Height           int     `json:"height"`
	BatchSize        int     `json:"batch_size"`
	NIter            int     `json:"n_iter"`
	Steps            int     `json:"steps"`
	ReturnGrid       bool    `json:"return_grid"`
	RestoreFaces     bool    `json:"restore_faces"`
	SendImages       bool    `json:"send_images"`
	SaveImages       bool    `json:"save_images"`
	DoNotSaveSamples bool    `json:"do_not_save_samples"`
	DoNotSaveGrid    bool    `json:"do_not_save_grid"`
	TaskId           string  `json:"task_id"`
}

type SuccessResponse struct {
	Images     []string `json:"images"`
	Parameters Data     `json:"parameters"`
	Info       string   `json:"info"`
}

func Txt2ImgRequest(prompt string) {
	var data Data

	data.Prompt = prompt
	//{
	//	"denoising_strength": 0,
	//	"prompt": "puppy dogs",
	//	"negative_prompt": "",
	//	"seed": -1,
	//	"batch_size": 5,
	//	"n_iter": 1,
	//	"steps": 5,
	//	"cfg_scale": 7,
	//	"width": 512,
	//	"height": 512,
	//	"restore_faces": false,
	//	"tiling": false,
	//	"sampler_index": "Euler"
	//}
	data.NegativePrompt = ""
	data.Seed = -1
	data.BatchSize = 5
	data.NIter = 1
	data.Steps = 5
	data.CfgScale = 7
	data.Width = 512
	data.Height = 512
	data.RestoreFaces = false
	data.SamplerIndex = "Euler"

	data.SendImages = true
	data.TaskId = utils.GenerateTaskID()

	url := "http://" + g.Cfg().GetString("sd_webui.ip") + "/sdapi/v1/txt2img"

	fmt.Println("url:", url)
	jsonData, err := json.Marshal(data)
	if err != nil {
		fmt.Println("Error:", err.Error())
		return
	}

	resp, err := utils.SubmitPost(url, jsonData)
	if err != nil {
		fmt.Println("Error:", err.Error())
		return
	}
	defer resp.Body.Close()

	fmt.Println("请求完了")

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		fmt.Println("Error:", err.Error())
		return
	}

	var response SuccessResponse

	err = json.Unmarshal(body, &response)
	if err != nil {
		fmt.Println("Error:", err.Error())
	}

	dir := utils.WeekTime(data.TaskId)

	for i, img := range response.Images {
		name := filepath.Join(dir, "img"+strconv.Itoa(i+1)+".png")
		err = utils.SaveEncodedImage(img, name)
		if err != nil {
			fmt.Println("Error:", err.Error())
			return
		}
		response.Images[i] = name
	}

	result, _ := json.Marshal(response)

	err = utils.SetTxt2ImgTaskToRedis(data.TaskId, string(result))
	if err != nil {
		fmt.Println("Error:", err.Error())
		return
	}

	fmt.Println("Task ID:", data.TaskId)
}
